from ultralytics import YOLO
import os
import cv2
import numpy as np

import CropSegDataset


save_dir = r'D:\lliujian\DeepLearning\Projects\LTTools\results'

# Load a model
model = YOLO("yolo11n-seg.pt")

input_img_path = r"D:\lliujian\DeepLearning\屏幕截图 2025-01-16 162043.png"
input_img = cv2.imdecode(np.fromfile(input_img_path, dtype=np.uint8), -1)
input_img = input_img[:, :, :3]
h, w, c = input_img.shape

# 推理单张图
crop_images, _, crop_positions = CropSegDataset.crop_seg_data(input_img, input_img[:, :, 0], 640, 640, 0.9)
output_mask = np.zeros((h, w))
for i, (inpt_im, inpt_pos) in enumerate(zip(crop_images, crop_positions)):
    results = model(inpt_im, retina_masks=True, imgsz=640)
    # results = model.predictor(input_img=input_img, retina_masks=True, imgsz=640)

    if results[0].masks is not None and len(results[0].masks) > 0:
        masks = results[0].masks.data.detach().cpu().numpy()
        for j, msk in enumerate((masks * 255).astype(np.uint8)):
            if j == 0:
                add_mask = msk
            else:
                add_mask += msk
        l, t, r, b = inpt_pos
        output_mask[t:b, l:r] += add_mask
        output_mask = np.clip(output_mask, a_min=0.0, a_max=1.0)

        # 保存裁剪图结果
        add_mask = np.clip(add_mask, a_min=0.0, a_max=1.0)
        cv2.imencode('.bmp', add_mask.astype(np.uint8) * 255)[1].tofile(
            os.path.join(save_dir, os.path.basename(input_img_path).replace('.png', f'_{i}.bmp')))
    else:
        continue

cv2.imencode('.bmp', output_mask.astype(np.uint8) * 255)[1].tofile(os.path.join(save_dir, os.path.basename(input_img_path).replace('.png', '.bmp')))
#
# contours, hierarchy = cv2.findContours(output_mask.astype(np.uint8) * 255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# draw_img = cv2.drawContours(input_img, contours, color='green', thickness=2)
# cv2.imencode('.bmp', draw_img.astype(np.uint8) * 255)[1].tofile(os.path.join(save_dir, os.path.basename(input_img_path).replace('.png', '_masked.bmp')))

